import re
import os
import argparse

"""
功能特点：
✅ 自动添加首行缩进

✅ 合并句子中间的错误换行

✅ 规范化段落格式

✅ 清理多余空行

✅ 支持中文标点识别

✅ 批量处理功能

✅ 自动编码检测

这个脚本应该能很好地解决你提到的小说格式问题！
"""
def format_novel(input_file, output_file=None, indent_spaces=4, encoding='utf-8'):
    """
    格式化小说文本
    
    参数:
    - input_file: 输入文件路径
    - output_file: 输出文件路径（默认为原文件名+_formatted）
    - indent_spaces: 缩进空格数（默认4个空格）
    - encoding: 文件编码（默认utf-8）
    """
    
    # 设置输出文件名
    if output_file is None:
        base_name = os.path.splitext(input_file)[0]
        output_file = f"{base_name}_formatted.txt"
    
    try:
        # 读取文件
        with open(input_file, 'r', encoding=encoding) as f:
            content = f.read()
    except UnicodeDecodeError:
        # 如果utf-8失败，尝试gbk编码
        try:
            with open(input_file, 'r', encoding='gbk') as f:
                content = f.read()
        except UnicodeDecodeError:
            print(f"错误：无法解码文件 {input_file}，请检查文件编码")
            return False
    
    # 处理文本格式
    formatted_content = process_content(content, indent_spaces)
    
    # 写入输出文件
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(formatted_content)
        print(f"格式化完成！输出文件：{output_file}")
        return True
    except Exception as e:
        print(f"写入文件时出错：{e}")
        return False

def process_content(content, indent_spaces=4):
    """
    处理文本内容的格式
    """
    # 1. 统一换行符为\n
    content = content.replace('\r\n', '\n').replace('\r', '\n')
    
    # 2. 合并错误的换行（句子中间换行）
    content = merge_broken_lines(content)

    content = is_novel_chapter_title(content)
    
    # 3. 规范化段落
    content = normalize_paragraphs(content)
    
    # 4. 添加首行缩进
    content = add_indentation(content, indent_spaces)
    
    # 5. 清理多余的空行
    content = clean_extra_newlines(content)
    
    return content

"""
判断是否为小说章节标题,如果为章节行前添加一行有20个星号符,的行则最后添加两个换行符
"""
def is_novel_chapter_title(content):
    """
    判断是否为小说章节标题，如果为章节行前添加一行有20个星号符，标题行的最后添加两个换行符
    """
    lines = content.split('\n')
    formatted_lines = []
    for line in lines:
        if is_chapter_title(line):
            formatted_lines.append('*' * 20)  # 添加20个星号符
            formatted_lines.append(line + '\n')  # 标题行后加一个换行
            formatted_lines.append('')  # 再加一个空行（即两个换行）
        else:
            formatted_lines.append(line)
    return '\n'.join(formatted_lines)

def is_chapter_title(line):
    """
    判断是否为小说章节标题
    """
    patterns = [
        r'^\s*第[\d一二三四五六七八九十百千万零两]+章.*',   # 第23章、第二十三章
        r'^\s*楔子.*',                             # 楔子
        r'^\s*前言.*',                             # 前言
        r'^\s*作者.*',                             # 作者
        r'^\s*（[\d一二三四五六七八九十百千万零两]+）.*',   # （二）、（23）、（第二十三）
    ]
    for pattern in patterns:
        if re.match(pattern, line.strip()):
            return True
    return False
    

def merge_broken_lines(content):
    """
    合并句子中间的错误换行
    """
    lines = content.split('\n')
    cleaned_lines = [line for line in lines if line.strip() != '']
    content = '\n'.join(cleaned_lines)

    # 定义句子结束标点（中文和英文）
    sentence_endings = r'[。！？!?]["」』】）)※]?'
    sentence_startings = r'^[「『【（"「『【（A-Z]'
    
    lines = content.split('\n')
    merged_lines = []
    current_paragraph = []
    
    for i, line in enumerate(lines):
        line = line.strip()
        
        if not line:  # 空行，段落分隔
            if current_paragraph:
                merged_lines.append(' '.join(current_paragraph))
                current_paragraph = []
            merged_lines.append('')
            continue
        
        # 去除每行中的空格
        line = re.sub(r' +','', line)

        # 检查是否是句子开头（新段落开始）
        if (not current_paragraph or 
            re.search(sentence_endings, current_paragraph[-1][-3:]) or  # 检查最后几个字符
            re.match(sentence_startings, line)):
            
            if current_paragraph:
                merged_lines.append(''.join(current_paragraph))
            current_paragraph = [line]
        else:
            # 合并到当前段落
            current_paragraph.append(line)
    
    # 添加最后一个段落
    if current_paragraph:
        merged_lines.append(' '.join(current_paragraph))
    
    return '\n'.join(merged_lines)

def normalize_paragraphs(content):
    """
    规范化段落格式，将段落内的空格替换为换行符
    """
    paragraphs = content.split('\n\n')
    normalized_paragraphs = []
    
    for para in paragraphs:
        if not para.strip():
            normalized_paragraphs.append('')
            continue
        
        # 将段落内的换行符替换为空格（合并成一段）
        para = ' '.join(para.split('\n'))
        
        # 清理多余的空格
        para = re.sub(r' +', ' ', para.strip())
        
        # 将段落内的空格替换为换行符
        para = para.replace(' ', '\n')
        
        normalized_paragraphs.append(para)
    
    return '\n\n'.join(normalized_paragraphs)

def add_indentation(content, indent_spaces=4):
    """
    添加首行缩进
    """
    indent = ' ' * indent_spaces
    lines = content.split('\n')
    formatted_lines = []
    
    for line in lines:
        if line.strip():  # 非空行
            # 检查是否已经是缩进格式（避免重复缩进）
            if not line.startswith((' ', '　', '\t')) and line.strip():
                formatted_lines.append(indent + line)
            else:
                formatted_lines.append(line)
        else:  # 空行
            formatted_lines.append('')
    
    return '\n'.join(formatted_lines)

def clean_extra_newlines(content):
    """
    清理多余的空行（最多保留2个空行）
    """
    # 将3个或更多连续空行替换为2个空行
    content = re.sub(r'\n\s*\n\s*\n+', '\n\n', content)
    return content.strip()

def batch_process_folder(folder_path, indent_spaces=4, encoding='utf-8'):
    """
    批量处理文件夹中的所有txt文件
    """
    if not os.path.exists(folder_path):
        print(f"错误：文件夹 {folder_path} 不存在")
        return
    
    txt_files = [f for f in os.listdir(folder_path) if f.endswith('.txt')]
    
    if not txt_files:
        print("文件夹中没有找到txt文件")
        return
    
    print(f"找到 {len(txt_files)} 个txt文件，开始批量处理...")
    
    success_count = 0
    for file_name in txt_files:
        input_file = os.path.join(folder_path, file_name)
        output_file = os.path.join(folder_path, f"formatted_{file_name}")
        
        print(f"正在处理：{file_name}")
        if format_novel(input_file, output_file, indent_spaces, encoding):
            success_count += 1
    
    print(f"批量处理完成！成功处理 {success_count}/{len(txt_files)} 个文件")

def main():
    parser = argparse.ArgumentParser(description='小说文本格式化工具')
    parser.add_argument('input', help='输入文件或文件夹路径')
    parser.add_argument('-o', '--output', help='输出文件路径')
    parser.add_argument('-i', '--indent', type=int, default=4, 
                       help='缩进空格数（默认：4）')
    parser.add_argument('-e', '--encoding', default='utf-8',
                       help='文件编码（默认：utf-8）')
    parser.add_argument('-b', '--batch', action='store_true',
                       help='批量处理文件夹中的所有txt文件')
    
    args = parser.parse_args()
    
    if args.batch:
        if os.path.isdir(args.input):
            batch_process_folder(args.input, args.indent, args.encoding)
        else:
            print("错误：批量处理模式需要输入文件夹路径")
    else:
        if os.path.isfile(args.input):
            format_novel(args.input, args.output, args.indent, args.encoding)
        else:
            print("错误：输入文件不存在")

if __name__ == "__main__":
    # 使用方法示例
    if len(os.sys.argv) == 1:
        print("小说格式化工具有两种使用方法：")
        print("\n1. 处理单个文件：")
        print("   python novel_formatter.py input.txt -o new.txt")
        print("\n2. 批量处理文件夹：")
        print("   python novel_formatter.py 文件夹路径 -b")
        print("\n3. 在代码中直接使用：")
        print("   format_novel('input.txt', 'output.txt', indent_spaces=4)")
    else:
        main()